# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/2/6 10:46 PM
# @Function:
class Runoff:
    """
    流域产流
    """

    def __init__(self):
        self.PE = None
        self.R = None
        self.RIM = None
        self.WDM = None
        self.A = None
        self.WMM = None
        self.EP = None
        self.ED = None
        self.EL = None
        self.EU = None
        self.P = None
        self.W = None
        self.WD = None
        self.WL = None
        self.WU = None
        self.WLM = None
        self.WUM = None
        self.IM = None
        self.B = None
        self.WM = None

    def SetParmameter(self, parameter):
        """
        设置参数
        Args:
            parameter:

        Returns:

        """
        self.WM = parameter.m_WM
        self.B = parameter.m_B
        self.IM = parameter.m_IM
        self.WUM = parameter.m_UM
        self.WLM = parameter.m_LM

    def SetState(self, state):
        self.WU = state.m_WU
        self.WL = state.m_WL
        self.WD = state.m_WD
        self.W = state.m_W
        self.P = state.m_P
        self.EU = state.m_EU
        self.EL = state.m_EL
        self.ED = state.m_ED
        self.EP = state.m_EP

    def calculate(self):
        """
        计算产流
        :return:
        """
        self.WMM = (1 + float(self.B)) / (1 - float(self.IM)) * float(self.WM)  # 包气带蓄水容量最大值，mm

        self.A = float(self.WMM) * (1 - pow(1 - float(self.W) / float(self.WM), 1 / (1 + float(self.B))))  # 初始土壤含水量最大值，mm

        self.PE = float(self.P) - float(self.EP)

        if self.PE <= 1e-5:  # 这里认为净雨量小于1e-5时即为小于等于0
            self.R = 0.0

            self.RIM = 0.0  # 计算不透水面积上的产流量
        else:
            if self.A + self.PE <= self.WMM:
                self.R = float(self.PE) + float(self.W) - float(self.WM) + float(self.WM) * pow(1 - (float(self.A) + float(self.PE)) / float(self.WMM), float(self.B) + 1)
            else:
                self.R = float(self.PE) - (float(self.WM) - float(self.W))

            self.RIM = float(self.PE) * float(self.IM)  # 计算不透水面积上的产流量

        # ========计算下一时段初土壤含水量========#
        self.WU = float(self.WU) + float(self.P) - float(self.EU) - float(self.R)

        self.WL = float(self.WL) - float(self.EL)

        self.WD = self.WD - self.ED

        if self.WD < 0:
            self.WD = 0  # 防止深层张力水蓄量小于0

        # 放置张力水蓄量超上限
        if self.WU > float(self.WUM):
            self.WL = float(self.WL) + float(self.WU) - float(self.WUM)
            self.WU = self.WUM

        if self.WL > float(self.WLM):
            self.WD = float(self.WD) + float(self.WL) - float(self.WLM)
            self.WL = self.WLM

        self.WDM = float(self.WM) - float(self.WUM) - float(self.WLM)  # 计算深层张力水容量

        if self.WD > self.WDM:
            self.WD = self.WDM

        # 计算土壤含水量
        self.W = float(self.WU) + float(self.WL) + float(self.WD)

    def UpdateState(self, state):
        state.m_WU = self.WU
        state.m_WL = self.WL
        state.m_WD = self.WD
        state.m_W = self.W
        state.m_R = self.R
        state.m_PE = self.PE
        state.m_RIM = self.RIM
